{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('../')\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from scipy import sparse\n",
    "from torch.utils.data import Dataset\n",
    "from tqdm import tqdm\n",
    "\n",
    "from MyPlot import *\n",
    "from utils import PieceWiseConst\n",
    "from ConstCofFVM import UniformFVM, BlockCofProblem\n",
    "from BaseTester import BaseTester\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class C1TestDs(Dataset):\n",
    "\tdef __init__(self, cofs, refs, dtype, device):\n",
    "\t\tself.cofs = torch.from_numpy(cofs).to(dtype)\n",
    "\t\tself.refs = torch.from_numpy(refs).to(dtype)\n",
    "\t\tself.dtype = dtype\n",
    "\t\tself.device = device\n",
    "\t\t\t\n",
    "\tdef __len__(self):\n",
    "\t\treturn len(self.cofs)\n",
    "\n",
    "\tdef __getitem__(self, index):\n",
    "\t\tcof = self.cofs[index]\n",
    "\n",
    "\t\tdata = cof.to(self.dtype).to(self.device)\n",
    "\n",
    "\t\tref = self.refs[index]\n",
    "\t\tref = ref.to(self.dtype).to(self.device)\n",
    "\n",
    "\t\treturn data[None, ...], ref, cof\n",
    "\n",
    "class C3TestDs(Dataset):\n",
    "\tdef __init__(self, cofs, refs, xx, yy, dtype, device):\n",
    "\t\tself.cofs = torch.from_numpy(cofs).to(dtype)\n",
    "\t\tself.refs = torch.from_numpy(refs).to(dtype)\n",
    "\t\tself.xx = torch.from_numpy(xx).to(dtype)\n",
    "\t\tself.yy = torch.from_numpy(yy).to(dtype)\n",
    "\n",
    "\t\tself.dtype = dtype\n",
    "\t\tself.device = device\n",
    "\t\t\t\n",
    "\tdef __len__(self):\n",
    "\t\treturn len(self.cofs)\n",
    "\n",
    "\tdef __getitem__(self, index):\n",
    "\t\tcof = self.cofs[index]\n",
    "\t\tdata = torch.stack([self.xx, self.yy, cof])\n",
    "\t\tdata = data.to(self.dtype).to(self.device)\n",
    "\n",
    "\t\tref = self.refs[index]\n",
    "\t\tref = ref.to(self.dtype).to(self.device)\n",
    "\n",
    "\t\treturn data, ref, cof\n",
    "\n",
    "class CofTester(BaseTester):\n",
    "\tdef __init__(self, **kwargs):\n",
    "\t\tsuper().__init__(**kwargs)\n",
    "\t\tself.mesh()\n",
    "\t\n",
    "\tdef init_test_ds(self, DataN, GridSize, area=((0, 0), (1, 1)), save_path=None, load_path=None):\n",
    "\t\tif load_path is None:\n",
    "\t\t\tcofs, refs = [], []\n",
    "\t\t\t# Randomly sample data\n",
    "\t\t\tfor t in [2, 3, 4]:\n",
    "\t\t\t\tfor _ in range(DataN):\n",
    "\t\t\t\t\tmu = np.random.uniform(0.1, 10, (t, t))\n",
    "\t\t\t\t\tpwc = PieceWiseConst(mu, area)\n",
    "\t\t\t\t\tcofs.append(pwc(self.xx, self.yy))\n",
    "\n",
    "\t\t\tsolver = UniformFVM(area, GridSize, GridSize, None)\n",
    "\t\t\tfor idx, cof in tqdm(enumerate(cofs)):\n",
    "\t\t\t\tproblem = BlockCofProblem(cof, GridSize, area)\n",
    "\t\t\t\tA = solver.get_A(problem)\n",
    "\t\t\t\tb = solver.get_B(problem)\n",
    "\t\t\t\tu = sparse.linalg.spsolve(A.tocsr(), b).reshape(GridSize, GridSize)\n",
    "\t\t\t\trefs.append(u)\n",
    "\n",
    "\t\t\trefs = np.stack(refs)\n",
    "\t\t\tcofs = np.stack(cofs)\n",
    "\n",
    "\t\t\tif not save_path is None:\n",
    "\t\t\t\tp = Path(f'{save_path}/{GridSize}')\n",
    "\t\t\t\tif not p.is_dir():\n",
    "\t\t\t\t\tp.mkdir(parents=True)\n",
    "\t\t\t\tnp.save(p/'cofs.npy', cofs)\n",
    "\t\t\t\tnp.save(p/'refs.npy', refs)\n",
    "\n",
    "\t\telse:\n",
    "\t\t\trefs = np.load(f'{load_path}/{GridSize}/refs.npy')\n",
    "\t\t\tcofs = np.load(f'{load_path}/{GridSize}/cofs.npy')\n",
    "\t\t\trefs = refs[-DataN:]\n",
    "\t\t\tcofs = cofs[-DataN:]\n",
    "\t\tprint(cofs.shape)\n",
    "\t\tprint(self.kwargs['net_kwargs']['in_channels'])\n",
    "\t\t\n",
    "\t\tmatch self.kwargs['net_kwargs']['in_channels']:\n",
    "\t\t\tcase 1:\n",
    "\t\t\t\ttest_ds = C1TestDs(cofs, refs, self.dtype, self.device)\n",
    "\t\t\tcase 3:\n",
    "\t\t\t\ttest_ds = C3TestDs(cofs, refs, self.xx, self.yy, self.dtype, self.device)\n",
    "\t\tself.ds = test_ds\n",
    "\t\t# return test_ds\n",
    "\t\n",
    "\tdef test(self, exp_name, DataN, best_or_last, save_path, load_path):\n",
    "\t\tself.load_kwargs(exp_name)\n",
    "\t\tself.init_test_ds(DataN, self.GridSize, self.area, save_path=save_path, load_path=load_path)\n",
    "\t\tself.load_ckpt(best_or_last, exp_name)\n",
    "\n",
    "\t\tdf = {\n",
    "\t\t\t'l2': []\n",
    "\t\t}\n",
    "\t\twith torch.no_grad():\n",
    "\t\t\tfor i, (data, ref, cof) in enumerate(self.ds):\n",
    "\t\t\t\tdata = data[None,...]\n",
    "\t\t\t\tpre = self.net(data)\n",
    "\n",
    "\t\t\t\tl2 = self.l2(pre, ref).item()\n",
    "\t\t\t\t# print(l2)\n",
    "\t\t\t\t# if l2 >= 9e-4:\n",
    "\t\t\t\t# \tcontinue\n",
    "\t\t\t\tdf['l2'].append(l2)\n",
    "\n",
    "\t\t\t\t# self.save_img(f\"{self.img_save_path}/{exp_name}/TestCase-{i}\", pre, ref, cof)\n",
    "\t\tp = Path(f\"{self.img_save_path}/{exp_name}\")\n",
    "\t\tif not p.is_dir():\n",
    "\t\t\tp.mkdir(parents=True)\n",
    "\t\tdf = pd.DataFrame(df)\n",
    "\t\tdf.to_csv(f\"{self.img_save_path}/{exp_name}/l2.csv\", index=False)\n",
    "\t\t\n",
    "\tdef save_img(self, path, pre, ans, cof):\n",
    "\t\tp = Path(path)\n",
    "\t\tif not p.is_dir():\n",
    "\t\t\tp.mkdir(parents=True)\n",
    "\t\tplt.rcParams['font.size'] = 18  # Default font size\n",
    "\n",
    "\t\tpre = pre.cpu().numpy().reshape(self.GridSize, self.GridSize)\n",
    "\t\tans = ans.cpu().numpy().reshape(self.GridSize, self.GridSize)\n",
    "\t\tcof = cof.cpu().numpy().reshape(self.GridSize, self.GridSize)\n",
    "\n",
    "\t\tsave_surf(path, pre, self.xx, self.yy, 'surf_pre')\n",
    "\t\tsave_surf(path, ans, self.xx, self.yy, 'surf_ref')\n",
    "\t\tsave_ctf(path, pre, ans, self.xx, self.yy)\n",
    "\t\tsave_contour(path, pre, ans, self.xx, self.yy, levels=None)\n",
    "\t\tsave_img_force(path, cof, 'cof', vmin=0.1, vmax=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "6000it [3:08:05,  1.88s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(6000, 96, 96)\n",
      "3\n"
     ]
    }
   ],
   "source": [
    "cof_tester = CofTester(\n",
    "    GridSize=96,\n",
    "    area=((0, 0), (1, 1)),\n",
    "    ckpt_save_path=f'model_save',\n",
    "    hyper_parameters_save_path = f'./hyper_parameters', \n",
    "    img_save_path = f'./images', \n",
    "    device='cuda',\n",
    "    dtype=torch.double,\n",
    "    # mus = np.random.uniform(0.1, 10, (10, 2, 2)),\n",
    "\t)\n",
    "\n",
    "exp_name = 'Conv_Double'\n",
    "DataN = 2000\n",
    "best_or_last = 'last'\n",
    "cof_tester.test(exp_name, DataN, best_or_last, save_path=None, load_path=None)\n",
    "# cof_tester.init_test_ds()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "# df = pd.read_csv('./JJQC3-96-ResBottleNeck-4#4#6#6#8-4#4#6#6#8-2-layer-max-reflect-jac-15-10000-5.csv')\n",
    "df = pd.read_csv('l2.csv')\n",
    "print(df.shape)\n",
    "fig = plt.figure(figsize=(10, 6))  # Set the figure size\n",
    "plt.rcParams['font.size'] = 18  # Default font size\n",
    "data = np.log10(df['l2'].values)\n",
    "sns.histplot(\n",
    "data,\n",
    "bins=50,  # Number of bins\n",
    "kde=True,  # Add KDE (Kernel Density Estimate)\n",
    "stat='probability',  # Normalize the histogram\n",
    "color='navy',  # Color of the bars\n",
    "edgecolor='black'  # Color of the edges\n",
    ")\n",
    "print(df['l2'].mean())\n",
    "\n",
    "# Customization\n",
    "plt.title(r'$\\mathcal{P}_\\boldsymbol{\\mu}$')  # Title of the plot\n",
    "plt.xlabel('$Log_{L_2}$')  # X-axis label\n",
    "plt.xticks(np.linspace(data.min(), data.max(), 7))\n",
    "plt.ylabel('Probability')  # Y-axis label\n",
    "plt.grid(True)  # Show grid lines\n",
    "plt.show()\n",
    "fig.savefig(f'test_all.png')\n",
    "plt.close()\n",
    "\n",
    "DataN = 2000\n",
    "for i in range(3):\n",
    "\tfig = plt.figure(figsize=(10, 6))  # Set the figure size\n",
    "\tplt.rcParams['font.size'] = 18  # Default font size\n",
    "\tdata = np.log10(df['l2'].values[i*DataN: (i+1)*DataN])\n",
    "\tsns.histplot(\n",
    "\tdata,\n",
    "\tbins=50,  # Number of bins\n",
    "\tkde=True,  # Add KDE (Kernel Density Estimate)\n",
    "\tstat='probability',  # Normalize the histogram\n",
    "\tcolor='crimson',  # Color of the bars\n",
    "\tedgecolor='black'  # Color of the edges\n",
    "\t)\n",
    "\tprint(df['l2'].mean())\n",
    "\n",
    "\t# Customization\n",
    "\tplt.title(f'$t={i+2}$')  # Title of the plot\n",
    "\tplt.xlabel('$Log_{L_2}$')  # X-axis label\n",
    "\tplt.xticks(np.linspace(data.min(), data.max(), 7))\n",
    "\tplt.ylabel('Probability')  # Y-axis label\n",
    "\tplt.grid(True)  # Show grid lines\n",
    "\tplt.show()\n",
    "\tfig.savefig(f'test_t={i+2}')\n",
    "\tplt.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
